import logging
import os

from llama_index.core import VectorStoreIndex, Settings, SummaryIndex
from llama_index.core.retrievers import RouterRetriever
from llama_index.core.schema import TextNode
from llama_index.core.selectors import LLMSingleSelector
from llama_index.core.tools import RetrieverTool
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.dashscope import DashScope

logging.basicConfig(level=logging.DEBUG)
print('初始化模型对象')
llm_dash = DashScope(
    model="qwen-turbo",  # 使用通义千问Max模型 qwen-max
    api_key=os.getenv("DASHSCOPE_API_KEY"),  # 从阿里云控制台获取
    temperature=0.3
)
embed_model_zh = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-zh-v1.5",  # 中文优化模型
    device="cpu"
)
# 配置全局设置
Settings.embed_model = embed_model_zh
Settings.llm = llm_dash

nodes1 = [
    TextNode(
        text="一起事件是指意外或恶意的事故，它有可能对我们的IT资产安全造成不良影响。",
        metadata={"department": "Security", "security_level": 2},
    ),
    TextNode(
        text="在高科技发展的行业中,好的IT资产通常会带来优质的回报,特别是在上海,很多高科技企业,上海更是一座美丽的城市,并且很多有爱心的人都养了宠物",
        metadata={"department": "Security", "security_level": 5},
    )
]

nodes2 = [
    TextNode(
        text="猫是很多人都喜欢养的宠物,特别是上海的一些IT人士",
        metadata={"department": "Security", "security_level": 2},
    ),
    TextNode(
        text="赵敏毕业后回到上海一直在家照顾父母,可以说是全职儿女了.",
        metadata={"department": "Security", "security_level": 5},
    )
]

# 修正1：直接构建索引并启用元数据存储
vector_index = VectorStoreIndex(nodes=nodes1)
vector_retriever = vector_index.as_retriever()

# 应用过滤器到检索器
summary_index = SummaryIndex(nodes=nodes2)
summary_retriever = summary_index.as_retriever()

vector_tool = RetrieverTool.from_defaults(
    retriever=vector_retriever,
    description="使用这个来回答有关上海的问题。"
)
summary_tool = RetrieverTool.from_defaults(
    retriever=summary_retriever,
    description="使用这个来回答有关宠物的问题。"
)

routerRetriever = RouterRetriever(
    selector=LLMSingleSelector.from_defaults(),
    retriever_tools=[
        vector_tool,
        summary_tool
    ]
)

# result = routerRetriever.retrieve("上海")
result = routerRetriever.retrieve("宠物")
print('检索结果:')
for item in result:
    print(f"==> {item.text} [元数据: {item.metadata}]")
